import random
import string

import easyocr
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image, ImageDraw, ImageFont


def generate_captcha(text_length=4, size=(150, 50), font_size=24):
    """Generate simple captcha image with text and noise"""
    # Generate random alphanumeric text
    chars = string.ascii_letters + string.digits
    captcha_text = "".join(random.choices(chars, k=text_length))

    # Create image canvas
    img = Image.new("RGB", size, (255, 255, 255))
    draw = ImageDraw.Draw(img)

    # Load font (use Windows default Arial font)
    try:
        font = ImageFont.truetype("arial.ttf", font_size)
    except IOError:
        font = ImageFont.load_default()

    # Draw text
    bbox = draw.textbbox((0, 0), captcha_text, font=font)  # 获取文本边界框(x0,y0,x1,y1)
    text_width = bbox[2] - bbox[0]  # 宽度 = 右边界 - 左边界
    text_height = bbox[3] - bbox[1]  # 高度 = 下边界 - 上边界
    x = (size[0] - text_width) / 2
    y = (size[1] - text_height) / 2
    draw.text((x, y), captcha_text, fill=(0, 0, 0), font=font)

    # Add noise lines
    for _ in range(3):
        x1 = random.randint(0, size[0])
        y1 = random.randint(0, size[1])
        x2 = random.randint(0, size[0])
        y2 = random.randint(0, size[1])
        draw.line([(x1, y1), (x2, y2)], fill=(128, 128, 128), width=1)

    return img, captcha_text


def ocr_recognize(image):
    """Recognize captcha text using EasyOCR"""
    img_np = np.array(image)
    reader = easyocr.Reader(["en"], gpu=False)  # Use CPU recognition
    result = reader.readtext(img_np, detail=0)
    return "".join(result).replace(" ", "")  # Clean up result


def main():
    # Generate captcha
    captcha_img, true_text = generate_captcha()

    # OCR recognition
    recognized_text = ocr_recognize(captcha_img)

    # Visualization
    plt.figure(figsize=(8, 4))
    plt.subplot(1, 2, 1)
    plt.imshow(captcha_img)
    plt.title(f"Generated Captcha\nTrue Text: {true_text}")

    plt.subplot(1, 2, 2)
    plt.imshow(captcha_img)
    plt.title(f"OCR Recognition\nResult: {recognized_text}")

    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    main()
